package cn.itcast.mapreduce.covid.partition;

import org.apache.hadoop.io.Text;
import org.apache.hadoop.mapreduce.Partitioner;

import java.util.HashMap;

public class StatePartitioner extends Partitioner<Text, Text> {

    //模拟美国各州的数据字典，实际中可以从redis进行读取加载 如果数据量不大 也可以创建数据集合保存
    public static HashMap<String, Integer> stateMap = new HashMap<String, Integer>();

    static {
        stateMap.put("Alabama", 0);
        stateMap.put("Alaska", 1);
        stateMap.put("Arkansas", 2);
        stateMap.put("California", 3);
        stateMap.put("Colorado", 4);
    }


    /**
     * todo 自定义分区器中分区规则的实现方法 只要getPartition返回的int一样  数据就会被分到同一个分区
     * 所谓的同一个分区指的是数据到同一个reducetask处理
     *
     * @param key           state：州
     * @param value         一行文本数据：2021-01-28,Dale,Alabama,01045,4235,92
     * @param numPartitions
     * @return
     */
    @Override
    public int getPartition(Text key, Text value, int numPartitions) {

        Integer code = stateMap.get(key.toString());

        if (code != null) {
            return code;
        }

        return 5;
    }

}
